#include "block_storage.h"
#include "pending_queue.h"
#include "proto.h"
#include "zookeeper_log.h"
#include "memory.h"
#include <linux/spinlock.h>
#include <linux/wait.h>
#include <linux/sched.h>
#include <linux/slab.h>

void check_pending_queue(block_storage_t * storage);
int keyToNode(long blockId, int* liveNodes, int* nodes);
int keyToNodeRead(long blockId, int* nodes);

block_storage_t * block_storage_init(const char* host, const int64_t client_id){
    clientid_t id;
    block_storage_t * storage = NULL;
    id.client_id = client_id;
    storage = calloc(1, sizeof(*storage));
    storage->client_id = client_id;
    storage->servers = zookeeper_init(host, 30000, &id, 0, 0, storage);
    storage->pending_queue = pending_queue_init();
    storage->xid = 0;
    storage->nonce = 0; //Should use random and better be long
    spin_lock_init(&storage->lock);
    return storage;
}

int get_outstanding_count(block_storage_t *storage){
    int ret = 0;
    unsigned long flags;
    spin_lock_irqsave(&storage->pending_queue->cond.lock, flags);
    ret = storage->pending_queue->count;
    spin_unlock_irqrestore(&storage->pending_queue->cond.lock, flags);
    //printk("outstanding count = %d\n", ret);
    return ret;
}



int32_t get_xid(block_storage_t * storage)
{
    int ret = 0;
    unsigned long flags;
    spin_lock_irqsave(&storage->lock, flags);
    ret = storage->xid++;
    spin_unlock_irqrestore(&storage->lock, flags);
    return ret;
}

int32_t get_nonce(block_storage_t * storage)
{
    int ret = 0;
    unsigned long flags;
    spin_lock_irqsave(&storage->lock, flags);
    ret = storage->nonce++;
    spin_unlock_irqrestore(&storage->lock, flags);
    return ret;
}

void init_xid(block_storage_t * storage, int32_t xid)
{
    unsigned long flags;
    spin_lock_irqsave(&storage->lock,flags);
    if(xid>=storage->xid)
        storage->xid = xid+1;
    spin_unlock_irqrestore(&storage->lock,flags);
}


void write_block(block_storage_t * storage,long blockId, char *buf, int buf_len,
        void_completion_t completion, const void *data){
    int xid = 0, nonce = 0;
    pending_write_t *p_write = NULL;
    int i;
    int nodes[SERVER_COUNT];
    int quorum[WRITE_QUORUM_SIZE];
    int list;
    unsigned long flags;
    //wait_event_interruptible_locked(storage->pending_queue->cond, storage->pending_queue->count<OUTSTANDING_LIMIT);
    if(storage->pending_queue->count>=OUTSTANDING_LIMIT){
	printk("queue is full\n");
	//spin_unlock_irqrestore(&storage->pending_queue->cond.lock,flags);
	return;
    }
    p_write = alloc_pending_write(blockId, xid, nonce, buf, buf_len, completion,(char*) data);
    if(p_write==NULL){
	printk("alloc_pending_write fails\n");
	//spin_unlock_irqrestore(&storage->pending_queue->cond.lock,flags);
	return;
    }
    spin_lock_irqsave(&storage->pending_queue->cond.lock,flags);
    xid = get_xid(storage);
    nonce = get_nonce(storage);
    p_write->xid = xid;
    p_write->nonce = nonce;
    add_to_tail(storage->pending_queue, p_write, nonce, PENDING_WRITE, zoo_recv_timeout(storage->servers));
    spin_unlock_irqrestore(&storage->pending_queue->cond.lock, flags);
    //return;
    
    get_live_servers(storage->servers, nodes);
    list = keyToNode(blockId, nodes, quorum);

    //for(i=0; i<WRITE_QUORUM_SIZE; i++){
	struct oarchive *oa;
        struct RequestHeader h = { .xid = xid, .nonce = nonce, .type = PREPAREDATA_OP};
        struct PrepareDataRequest req;
        int rc;

        oa = create_buffer_oarchive();
	if(oa==NULL){
	    printk("create oa failed\n");
	    return;
	}
        req.clientId = storage->client_id;
        req.count = list; //WRITE_QUORUM_SIZE;
        req.data.buff = buf;
        req.data.len = buf_len;
        rc = serialize_RequestHeader(oa, "header", &h);
        rc = rc < 0 ? rc : serialize_PrepareDataRequest(oa, "req", &req);
        send_bytes_to_server(storage->servers, get_buffer(oa),
            get_buffer_len(oa), quorum[0]);
        if(DEBUG_SHARD)
        printk("sending prepare nonce =%d for blockId [%ld] to %d\n",h.nonce,blockId, list);
        close_buffer_oarchive(&oa, 0);

    //}
}

void read_block(block_storage_t * storage, long blockId,
	 data_completion_t completion, const void *data){
    int nonce = 0;
    pending_read_t *p_read = NULL;
    int rc;
    int nodes[SERVER_COUNT];
    unsigned long flags;

    //debug
    /*char *ret = kmalloc(4096, GFP_KERNEL);
    if(ret==NULL){
	printk("alloc read return failed\n");
	return;
    }
    completion(0, ret, 4096, data);*/
    //return;
    //debug end

    //wait_event_interruptible_locked(storage->pending_queue->cond, storage->pending_queue->count<OUTSTANDING_LIMIT);
    p_read = alloc_pending_read(blockId, 0, completion,(char*) data);
    if(p_read==NULL){
        printk("alloc_pending_read fails\n");
        //spin_unlock_irqrestore(&storage->pending_queue->cond.lock, flags);
        return;
    }

    spin_lock_irqsave(&storage->pending_queue->cond.lock, flags);
    nonce = get_nonce(storage);
    p_read->nonce = nonce;
    add_to_tail(storage->pending_queue, p_read, nonce, PENDING_READ, zoo_recv_timeout(storage->servers));
    spin_unlock_irqrestore(&storage->pending_queue->cond.lock, flags);

    struct oarchive *oa;
    struct RequestHeader h = { .xid = -1, .nonce = nonce, .type = BOLOSKYGETDATA_OP};
    struct BoloskyGetDataRequest req;

    oa = create_buffer_oarchive();
    if(oa==NULL){
        printk("create oa failed\n");
        return;
    }

    req.blockId = blockId;
    get_live_servers(storage->servers, nodes);
    req.serverId = keyToNodeRead(blockId, nodes);
    p_read->serverId = req.serverId; 
    rc = serialize_RequestHeader(oa, "header", &h);
    rc = rc < 0 ? rc : serialize_BoloskyGetDataRequest(oa, "req", &req);

    send_bytes_to_server(storage->servers, get_buffer(oa),
            get_buffer_len(oa), SERVER_COUNT-1);
    if(DEBUG_SHARD)
    printk("sending read nonce=%d for blockId [%ld] to %d\n",h.nonce, blockId, SERVER_COUNT-1);
    close_buffer_oarchive(&oa, 0);


    //Debug
    /*spin_lock(&storage->pending_queue->cond.lock);
    
    p_read = get_by_nonce(storage->pending_queue, nonce);
    spin_unlock(&storage->pending_queue->cond.lock);
    
    if(p_read==NULL){
	printk("get_by_nonce returns NULL\n");
	return;
    }
    p_read->buf_len = 4096;
    p_read->buf = kmalloc(4996, GFP_KERNEL);
    if(p_read->buf==NULL){
	printk("malloc failed\n");
	return;
    }

    p_read->completion(0, p_read->buf, p_read->buf_len, p_read->data);

    spin_lock(&storage->pending_queue->cond.lock);
    remove_first(storage->pending_queue);
    spin_unlock(&storage->pending_queue->cond.lock);
    free(p_read);*/

    //End debug
}

void handle_response(block_storage_t * storage, struct ReplyHeader *hdr, struct iarchive *ia, int sid){
    pending_write_t *p_write = NULL;
    pending_read_t *p_read = NULL;
    struct BoloskyGetDataResponse res;
    //unsigned long flags;
    switch(hdr->type){
	case 3: //PREPAREDATA
	    if(DEBUG_SHARD)
	    printk("receive prepare reply nonce=%d, sid=%d\n", hdr->nonce, sid);
	    //spin_lock(&storage->pending_queue->cond.lock);
	    p_write = get_by_nonce(storage->pending_queue, hdr->nonce);
	    //spin_unlock(&storage->pending_queue->cond.lock);
	    if(p_write == NULL)
		break;
	    if(pending_write_data_ready(p_write)==0){
		pending_write_add_prepare_reply(p_write, sid);
		if(pending_write_data_ready(p_write)==1)
		    check_pending_queue(storage);
	    }
	    break;
	case 2: //SETDATA
	    if(DEBUG_SHARD)
	    printk("receive setdata response nonce=%d, sid=%d\n", hdr->nonce, sid);
	    //spin_lock(&storage->pending_queue->cond.lock);
            p_write = get_by_nonce(storage->pending_queue, hdr->nonce);
	    //spin_unlock(&storage->pending_queue->cond.lock);
            if(p_write == NULL)
                break;
            pending_write_add_write_reply(p_write, sid);
	    if(pending_write_ready(p_write))
		check_pending_queue(storage);
	    break;
	case 4: //BOLOSKYREAD
	    if(DEBUG_SHARD)
	    printk("receive read response nonce=%d, sid=%d\n", hdr->nonce, sid);
	    //spin_lock(&storage->pending_queue->cond.lock);
            p_read = get_by_nonce(storage->pending_queue, hdr->nonce);
	    //spin_unlock(&storage->pending_queue->cond.lock);
            if(p_read == NULL)
                break;
            deserialize_BoloskyGetDataResponse(ia, "reply", &res);
	    if(res.data.len==0){
		//This is reading an invalid block. Try read from another replica
		if(DEBUG_SHARD)
		    printk("An invalid read, retry\n");
		pending_read_timeout(storage, p_read);
	    }
	    else{
	    	p_read->buf_len = res.data.len;
	    	p_read->buf = kmalloc(res.data.len, GFP_KERNEL);
		if(p_read->buf==NULL){
		    printk("malloc returns NULL\n");
		    return;
		}
		memcpy(p_read->buf, res.data.buff, res.data.len);//res.data.buff;

            	pending_read_add_reply(p_read, sid);
            	if(pending_read_ready(p_read))
                    check_pending_queue(storage);
	    }
            deallocate_BoloskyGetDataResponse(&res);
            break;
	default:
	    printk("Unknow reply %d\n", hdr->type);
    }
}

static char set_data_buf[1];
void check_pending_queue(block_storage_t * storage){
    int i, type;
    void *p;
    int finish_count = 0;
    int count=0;
    unsigned long flags;
    spin_lock_irqsave(&storage->pending_queue->cond.lock, flags);
    count=storage->pending_queue->count;
    spin_unlock_irqrestore(&storage->pending_queue->cond.lock, flags);
    //printk("check_pending_queue\n");
    //First check if any write request is prepared
    for(i=0; i<count; i++){
	//spin_lock(&storage->pending_queue->cond.lock);
	p = get_by_index(storage->pending_queue, i, &type);
	//spin_unlock(&storage->pending_queue->cond.lock);
	if(p==NULL){
	    printk("get_by_index returns NULL\n");
	    break;
	}
	if(type == PENDING_WRITE){
	    pending_write_t *p_write = p;
	    if(pending_write_data_ready(p_write)==1){
		if(p_write->write_sent==0){
		    struct oarchive *oa;
                    struct RequestHeader h = { .xid = p_write->xid, .nonce = p_write->nonce, .type = SETDATA_OP};
                    struct SetDataRequest req;
                    int rc;

                    oa = create_buffer_oarchive();
		    if(oa==NULL){
			printk("create oa failed\n");
			//spin_unlock(&storage->pending_queue->cond.lock);
			return;
		    }
                    req.blockId = p_write->blockId;
                    req.data.buff = set_data_buf;
                    req.data.len = 1;
                    rc = serialize_RequestHeader(oa, "header", &h);
                    rc = rc < 0 ? rc : serialize_SetDataRequest(oa, "req", &req);

                    send_bytes_to_server(storage->servers, get_buffer(oa),
                        get_buffer_len(oa), 2);
		    if(DEBUG_SHARD)
                    printk("sending setdata xid=%d nonce=%d for blockId [%lld] to %d\n",h.xid, h.nonce,req.blockId,SERVER_COUNT-1);
                    close_buffer_oarchive(&oa, 0);
		    p_write->write_sent = 1;
		}
	    }
	    else
		break;
	}
    }

    //Second check if any request is ready to return to user
    for(i=0; i<count; i++){
	//spin_lock(&storage->pending_queue->cond.lock);
        p = get_by_index(storage->pending_queue, i, &type);
	//spin_unlock(&storage->pending_queue->cond.lock);
	if(p==NULL){
	    printk("get_by_index returns NULL\n");
	    break;
	}
        if(type == PENDING_WRITE){
            pending_write_t *p_write = p;
            if(pending_write_ready(p_write)==1){
		if(DEBUG_SHARD)
		printk("finish write block %ld, nonce %d, xid %d\n", p_write->blockId, p_write->nonce, p_write->xid);
		//spin_unlock(&storage->pending_queue->cond.lock);
		p_write->completion(0, p_write->data);
		//spin_lock(&storage->pending_queue->cond.lock);
		dealloc_pending_write(p_write);
		free(p_write);
		finish_count++;
            }
            else
                break;
        }
	else{
	    pending_read_t *p_read = p;
	    if(pending_read_ready(p_read)==1){
		if(DEBUG_SHARD)
		printk("finish read block %ld, nonce %d\n", p_read->blockId, p_read->nonce);
		//spin_unlock(&storage->pending_queue->cond.lock);
		p_read->completion(0, p_read->buf, p_read->buf_len, p_read->data);
		//spin_lock(&storage->pending_queue->cond.lock);
		//dealloc_pending_read(p_read);
                free(p_read);
		finish_count++;
	    }
	    else
		break;
	}
    }
    spin_lock_irqsave(&storage->pending_queue->cond.lock, flags);
    for(i=0; i<finish_count; i++)
	(remove_first(storage->pending_queue));
    /*if(finish_count>0 && storage->pending_queue->count<OUTSTANDING_LIMIT){
	wake_up_locked(&(storage->pending_queue->cond));
    }*/
    spin_unlock_irqrestore(&storage->pending_queue->cond.lock, flags);
}

void close(block_storage_t * storage){
    zookeeper_close(storage->servers, get_xid(storage));
    free(storage);
}

int keyToNode(long blockId, int* nodes, int* quorum){
    int shard = blockId%SERVER_COUNT;
    int index = 0;
    int i;
    int ret = 0;
    for(i=0; i<SERVER_COUNT; i++){
        int tmp = (shard - i + SERVER_COUNT)%SERVER_COUNT;
        if(nodes[tmp]==1){
            quorum[index++]=tmp;
            ret=ret*10+(tmp+1);
        }
        if(index==WRITE_QUORUM_SIZE)
            break;
    }
    if(index<2){//!=WRITE_QUORUM_SIZE){
        printk("Not enough servers: %d\n", index-1);
        return -1;
    }
    return ret;
}


int keyToNodeRead(long blockId, int* nodes){
    int shard = blockId%SERVER_COUNT;
    int i;
    for(i=0; i<SERVER_COUNT; i++){
        int tmp = (shard - i + SERVER_COUNT)%SERVER_COUNT;
        if(nodes[tmp]==1)
            return tmp;
    }
    return -1;
}


